import numpy as np
from utilities import *
from UItools import *
class PartData:
    ELEMENT_TYPE_MAPPING = {
        'S3': (vtk.vtkTriangle, 3),
        'S6': (vtk.vtkTriangle, 3),  # 二阶三角形, 退化表示
        'S4': (vtk.vtkQuad, 4),
        'S8': (vtk.vtkQuad, 4),      # 二阶四边形
        'C3D8': (vtk.vtkHexahedron, 8),
        'C3D20': (vtk.vtkHexahedron, 8),  # 二阶六面体
        'C3D4': (vtk.vtkTetra, 4),
        'C3D10': (vtk.vtkTetra, 4),      # 二阶四面体
        'C3D6': (vtk.vtkWedge, 6),
        'C3D15': (vtk.vtkWedge, 6)        # 二阶楔形体
    }
    def __init__(self, node_arr=None, element_dict:dict=None, node_id_arr=None):
        
        self.color_map = None
        
        # VTK 数据
        self.points = vtk.vtkPoints()
        self.scalars = vtk.vtkFloatArray()
        self.scalars.SetName("R")
        self.umesh = vtk.vtkUnstructuredGrid()
        self.mapper = vtk.vtkDataSetMapper()
        self.actor = vtk.vtkActor()
        self.actor.GetProperty().SetEdgeColor(0, 0, 0)
        
        if not element_dict is None:
            if type(node_id_arr) == dict:
                self._load_data(node_arr, element_dict, node_id_arr)
            else:
                self.load_data(node_arr, element_dict, node_id_arr)

    def load_data(self, node_arr, element_dict, node_id_arr=None):
        self.node_id_arr = node_id_arr or range(1, 1 + len(node_arr))
        node_id_to_index = {node_id: idx for idx, node_id in enumerate(self.node_id_arr)}
        self._load_data(node_arr, element_dict, node_id_to_index)

    def _load_data(self, node_arr:np.ndarray, 
                   element_dict:Dict[str, np.ndarray], 
                   node_id_to_index:Dict[int, int]):
        
        node_arr = np.ascontiguousarray(node_arr, dtype=np.float32)
        vtk_points = vtk_numpy.numpy_to_vtk(node_arr, deep=False)
        self.points.SetData(vtk_points)
        self.umesh.SetPoints(self.points)
        self._node_arr = node_arr

        for elem_type, elem_arr in element_dict.items():
            vtk_cell_class, num_nodes = self._get_vtk_cell_info(elem_type)
            if vtk_cell_class is None:
                continue
            
            for elem in elem_arr:
                cell:vtk.vtkQuad = vtk_cell_class()  
                for i in range(num_nodes):  
                    nidx  = node_id_to_index[elem[i]]
                    cell.GetPointIds().SetId(i, nidx)
                
                self.umesh.InsertNextCell(cell.GetCellType(), cell.GetPointIds())#
        
        self.actor.GetProperty().SetEdgeColor(0, 0, 0)

    def load_shared(self, part:"PartData"):
        "基于现有PartData引用数据，无需重新构建umesh"
        self.points = part.points
        self.umesh.ShallowCopy(part.umesh)

    def _get_vtk_cell_info(self, elem_type_name):
        """根据单元类型名称返回VTK类和所需节点数"""
        for key, (vtk_class, num_nodes) in self.ELEMENT_TYPE_MAPPING.items():
            if key in elem_type_name:  # 部分匹配
                return vtk_class, num_nodes
        return None, 0
    
    def set_cmap(self, color_map):
        self.color_map = color_map
        
    def only_geom(self):
        self.mapper.SetInputData(self.umesh)
        self.mapper.ScalarVisibilityOff()  # 禁用标量映射
        self.actor.GetProperty().SetColor(0, 0.5, 1)  # 设置为纯色
        self.actor.SetMapper(self.mapper)
        
    def update_contour(self, response_arr, mode=0):
        """ mode = 1: cell mode
            mode = 0: point mode """
        if response_arr is None or len(response_arr) == 0:
            return
        vtk_array = vtk_numpy.numpy_to_vtk(
            response_arr.astype(np.float32).ravel(), 
            deep=False
            )
        vtk_array.SetName("Rv")
        if mode == 1:
            if len(response_arr) != self.umesh.GetNumberOfCells():
                raise ValueError("单元数据长度不匹配")
            self.umesh.GetCellData().SetScalars(vtk_array)
            self.mapper.SetScalarModeToUseCellData()
        else:
            if len(response_arr) != self.umesh.GetNumberOfPoints():
                raise ValueError("点数据长度不匹配")
            self.umesh.GetPointData().SetScalars(vtk_array)
            self.mapper.SetScalarModeToUsePointData()
        self.scalars = vtk_array  # 更新引用
        self.enable_contour()
        
    def enable_contour(self):
        if self.scalars.GetSize() == 0 or self.color_map is None:return
        self.mapper.ScalarVisibilityOn()
        # 配置Mapper和Actor
        self.mapper.SetInputData(self.umesh)
        self.mapper.SetLookupTable(self.color_map) # 绑定共享的颜色表
        self.mapper.UseLookupTableScalarRangeOn()  # 使用颜色表的范围
        self.actor.SetMapper(self.mapper)
    
    def get_min_max_response(self):
        """返回当前部件的响应值范围"""
        if self.scalars.GetSize() == 0:
            return 0. , 0.
        return np.min(self.scalars), np.max(self.scalars)
    
    def set_color_map_range(self, min_val, max_val):
        """设置颜色映射范围"""
        self.mapper.SetScalarRange(min_val, max_val)
    
    def set_visibility(self, visible):
        """设置可见性"""
        self.actor.SetVisibility(visible)
    
    def set_edge_visibility(self, visible):
        """设置网格边线可见性"""
        self.actor.GetProperty().SetEdgeVisibility(visible)